"""
LangChain 股票分析智能体协作系统 - 主程序
展示 LangChain 的核心功能和最佳实践
"""
from agent_flow import LangChainAgentFlow
from dotenv import load_dotenv
from colorama import init, Fore, Style
import json
import os
import sys
from typing import Dict, Any


def display_langchain_features():
    """展示 LangChain 功能特性"""
    print(Fore.CYAN + Style.BRIGHT + "\n🔥 === LangChain 功能展示 ===")
    features = [
        "🔗 Runnable 链式调用 - 构建复杂的处理流程",
        "🤖 多智能体协作 - Agent 之间的数据传递和协作",
        "📝 Prompt Templates - 结构化的提示词管理",
        "🛠️ Custom Tools - 自定义工具集成",
        "📊 Output Parsers - 结构化输出解析",
        "⚡ Async Support - 异步处理能力",
        "🔄 Error Handling - 完善的错误处理机制",
        "📈 Workflow Orchestration - 工作流编排"
    ]
    
    for feature in features:
        print(Fore.YELLOW + f"  {feature}")
    
    print(Fore.CYAN + Style.BRIGHT + "\n本项目展示了以上所有核心功能的实际应用！")


def check_langchain_environment():
    """检查 LangChain 环境配置"""
    print(Fore.CYAN + "\n🔍 检查 LangChain 环境...")
    
    try:
        # 检查 LangChain 相关包
        import langchain_core
        import langchain_openai
        print(Fore.GREEN + "✅ LangChain Core 已安装")
        print(Fore.GREEN + "✅ LangChain OpenAI 已安装")
        
        # 检查配置文件
        if os.path.exists("config.json"):
            print(Fore.GREEN + "✅ 配置文件存在")
        else:
            print(Fore.RED + "❌ 配置文件缺失")
            return False
        
        # 检查环境变量
        load_dotenv()
        
        with open("config.json", "r", encoding="utf-8") as f:
            config = json.load(f)
        
        # 检查智能体配置
        agents_config = config.get("agents", {})
        for agent_name, agent_info in agents_config.items():
            api_key = agent_info.get("app_key")
            if api_key and not api_key.startswith("your_"):
                print(Fore.GREEN + f"✅ {agent_name} API Key 已配置")
            else:
                print(Fore.YELLOW + f"⚠️ {agent_name} API Key 需要配置")
        
        # 检查 Tushare Token
        tushare_token = config.get("tushare_token")
        if tushare_token and not tushare_token.startswith("your_"):
            print(Fore.GREEN + "✅ Tushare Token 已配置")
        else:
            print(Fore.YELLOW + "⚠️ Tushare Token 需要配置")
        
        return True
        
    except ImportError as e:
        print(Fore.RED + f"❌ LangChain 包缺失: {e}")
        print(Fore.YELLOW + "请运行: pip install -r requirements.txt")
        return False
    except Exception as e:
        print(Fore.RED + f"❌ 环境检查失败: {e}")
        return False


def interactive_demo():
    """交互式演示"""
    print(Fore.CYAN + Style.BRIGHT + "\n🎮 === 交互式演示模式 ===")
    
    while True:
        print(Fore.CYAN + "\n请选择演示选项:")
        print("1. 🚀 运行完整的智能体协作流程")
        print("2. 📊 查看智能体信息")
        print("3. 🔧 测试单个智能体")
        print("4. 📈 分析自定义股票")
        print("5. 🏃 退出演示")
        
        choice = input(Fore.GREEN + "\n请输入选项 (1-5): ").strip()
        
        if choice == "1":
            run_full_workflow()
        elif choice == "2":
            show_agent_info()
        elif choice == "3":
            test_individual_agent()
        elif choice == "4":
            analyze_custom_stock()
        elif choice == "5":
            print(Fore.CYAN + "👋 感谢使用 LangChain 智能体协作系统！")
            break
        else:
            print(Fore.RED + "❌ 无效选项，请重新选择")


def run_full_workflow():
    """运行完整工作流"""
    print(Fore.CYAN + "\n🚀 启动完整的 LangChain 智能体协作流程...")
    
    try:
        agent_flow = LangChainAgentFlow()
        stock_code = "600519.SH"  # 贵州茅台
        print(Fore.GREEN + f"📈 分析股票: {stock_code} (贵州茅台)")
        agent_flow.run(stock_code)
        
    except Exception as e:
        print(Fore.RED + f"❌ 工作流执行失败: {e}")


def show_agent_info():
    """显示智能体信息"""
    print(Fore.CYAN + "\n📊 LangChain 智能体信息...")
    
    try:
        agent_flow = LangChainAgentFlow()
        agent_flow.get_agent_info()
        
    except Exception as e:
        print(Fore.RED + f"❌ 获取智能体信息失败: {e}")


def test_individual_agent():
    """测试单个智能体"""
    print(Fore.CYAN + "\n🔧 单个智能体测试...")
    
    agents = {
        "1": "数据收集智能体",
        "2": "数据分析智能体", 
        "3": "投资决策智能体"
    }
    
    print("请选择要测试的智能体:")
    for key, name in agents.items():
        print(f"{key}. {name}")
    
    choice = input(Fore.GREEN + "请输入选项 (1-3): ").strip()
    
    if choice not in agents:
        print(Fore.RED + "❌ 无效选项")
        return
    
    try:
        agent_flow = LangChainAgentFlow()
        stock_code = "600519.SH"
        
        if choice == "1":
            print(Fore.YELLOW + "🔍 测试数据收集智能体...")
            result = agent_flow.data_collector.execute(stock_code)
            print(Fore.WHITE + f"结果: {result[:300]}...")
            
        elif choice == "2":
            print(Fore.YELLOW + "📊 测试数据分析智能体...")
            # 使用模拟数据
            mock_data = f"股票 {stock_code} 的模拟数据报告..."
            result = agent_flow.data_analyst.execute(mock_data)
            print(Fore.WHITE + f"结果: {result[:300]}...")
            
        elif choice == "3":
            print(Fore.YELLOW + "🎯 测试投资决策智能体...")
            # 使用模拟分析报告
            mock_analysis = f"股票 {stock_code} 的模拟分析报告..."
            result = agent_flow.decision_maker.execute(mock_analysis)
            print(Fore.WHITE + f"结果: {result[:300]}...")
            
    except Exception as e:
        print(Fore.RED + f"❌ 智能体测试失败: {e}")


def analyze_custom_stock():
    """分析自定义股票"""
    print(Fore.CYAN + "\n📈 自定义股票分析...")
    
    stock_code = input(Fore.GREEN + "请输入股票代码 (例如: 000001.SZ): ").strip()
    
    if not stock_code:
        print(Fore.RED + "❌ 股票代码不能为空")
        return
    
    try:
        agent_flow = LangChainAgentFlow()
        agent_flow.run(stock_code)
        
    except Exception as e:
        print(Fore.RED + f"❌ 股票分析失败: {e}")


def main():
    """主函数"""
    init(autoreset=True)
    
    print(Fore.CYAN + Style.BRIGHT + "=" * 60)
    print(Fore.CYAN + Style.BRIGHT + "🤖 LangChain 股票分析智能体协作系统")
    print(Fore.CYAN + Style.BRIGHT + "=" * 60)
    
    # 显示 LangChain 功能特性
    display_langchain_features()
    
    # 检查环境
    if not check_langchain_environment():
        print(Fore.RED + "\n❌ 环境检查失败，请检查配置后重试")
        return
    
    print(Fore.GREEN + "\n✅ 环境检查通过，系统准备就绪！")
    
    # 询问运行模式
    print(Fore.CYAN + "\n请选择运行模式:")
    print("1. 🚀 快速演示 (自动运行)")
    print("2. 🎮 交互式演示")
    
    mode = input(Fore.GREEN + "请输入选项 (1-2): ").strip()
    
    if mode == "1":
        print(Fore.CYAN + "\n🚀 启动快速演示模式...")
        run_full_workflow()
    elif mode == "2":
        interactive_demo()
    else:
        print(Fore.YELLOW + "使用默认快速演示模式...")
        run_full_workflow()
    
    print(Fore.CYAN + Style.BRIGHT + "\n🎉 感谢使用 LangChain 智能体协作系统！")


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print(Fore.YELLOW + "\n\n👋 用户中断，程序退出")
    except Exception as e:
        print(Fore.RED + f"\n❌ 程序异常退出: {e}")
        sys.exit(1)